import tensorflow as tf
import numpy as np
import math
import os
import matplotlib.pyplot as plt
from pathlib import Path
from IPython import display
import glob
from PIL import Image
# from src import waymo_training_math_tools as math_tool
from src import waymo_training_plot_tools as plot_tool
from src import waymo_scenario_classes as class_tool
from src.waymo_transfer_training_to_tf import *
from src.waymo_training_math_tools import *
data_set='/home/mdmp/prpproject/WOMD_Dataset/training_dataset/uncompressed_scenario_training_training.tfrecord-00166-of-01000'
waymo_dataset_list = []

waymo_dataset_list.append(data_set)

# print("====================LOAD THE DATASET=============================")

# scenarios_get_list = [0]
# scenarios_summary= process_waymo_training_dataset(data_set, scenarios_get_list)



# time_index_list = scenarios_summary.scenarios_list[0].timestep_list

# print("====================ONE SCENARIO(class Motion_training_scenario)=============================")
# print("time_index of this scenario is(9s in total):")
# print(time_index_list)

# scenario = scenarios_summary.scenarios_list[0]

# print("real timestep in seconds:")
# print(scenario.timestep_seconds_list)

# states_index_list = time_index_list   
# scenario_list = scenarios_summary.scenarios_list    
# print("Number of scenario we have loaded:")
# print(len(scenario_list)) 

# print(len(scenario.map_features_list))


# print("All objects we can use:")
# print(dir(scenario))

# print("current_time_index is :")
# print(scenario.current_time_index)



# print("====================AGENTS LIST=============================")

# agents_list = scenario.agents_list
# print("num of all agents:")
# print(len(agents_list))

# print("====================FOR ONE AGENT(use sdc vehicle as an example)( class:Motion_training_agent)=============================")
# sdc_index = scenario.sdc_track_index
# sdc_vehicle = agents_list[sdc_index]    # 自车索引

# print("objects we can use in one agent:")
# print(dir(sdc_vehicle)) # 自车

# # inex = math_tool.Id_to_Index(agents_list,787.00)
# # print("wo zai na")
# # print(inex)

# print("Use sdc vehicle as an example")
# print("agent_id is :")
# print(sdc_vehicle.agent_id)

# sdc_vehicle_type = sdc_vehicle.agent_type
# print("sdc vehicle's type(of course a vehicle),(set sdc_vehicle as -1, unset = 0, vehicle = 1, pedestrian = 2, cyclist = 3, other = 4):")
# print(sdc_vehicle_type)

# print("====================STATES OF ONE AGENT=============================")
# states_list = sdc_vehicle.states_list
# print("we can use object 'states_list' to get all states of one agent in 91 timesteps, which is stored in class 'Agent_states'")
# print("num of states of one agent(actually the num of timestep now):")
# print(len(states_list))


# states_array = sdc_vehicle.states_array
# print("Directly get state array(nparray) of one agent, states[0] is timestep index; states[1] is all state of this agent")
# print(states_array.shape)
# print(type(states_array))
# print(type(states_array[0]))
# print(type(states_array[1]))
# print(states_array[1].shape)
# print("jiaodushi:")
# print(states_array[1][9])
# print(agents_list[math_tool.Id_to_Index(agents_list,3179)].states_array[1][9])
# print("All 14 states:")
# print("(state.timestep, state.agent_id, state.agent_type, state.center_x, state.center_y, state.center_z, state.length, state.width, state.height, state.heading, state.velocity_x, state.velocity_y, state.speed, state.valid)")

# states_mask_list = sdc_vehicle.states_mask_list
# print("state validation signs list")
# print(states_mask_list)


# print("====================GET ONE STATE(class: Agent_state)=============================")
# print("object we can use in class 'Agent_states':")
# cur_state_sdc = states_list[scenario.current_time_index]
# print(dir(cur_state_sdc))
# print("validation sign of state at current time:")
# print(cur_state_sdc.valid)
# agent_x = cur_state_sdc.center_x
# agent_y = cur_state_sdc.center_y
# agent_width = cur_state_sdc.width
# agent_length = cur_state_sdc.length
# print((agent_x,agent_y))

# agent_id = sdc_vehicle.agent_id

# xs = [o.center_x for o in states_list[9:11]]
# print(xs)

dataloader = Dataloader(data_set)



length_l = dataloader.get_ego_vehicle_length_list(80,90)
print(length_l)

pos = dataloader.get_ego_vehicle_pos_list(80,90)
print(pos)
ego_agent_id = dataloader.get_ego_vehicle_id()


front_car_list = dataloader.get_front_vehicle_id_list(ego_agent_id,0,90)
print(front_car_list)
back_car_list = dataloader.get_back_vehicle_id_list(ego_agent_id,0,90)
print(back_car_list)


acc, the =  dataloader.get_ego_vehicle_acc_and_acctheta_list(0,90)
print(acc)
# acc_1, the_1 = dataloader.get_agent_acc_and_acctheta_list_byId(front_car_list[0],0,90)
# print(acc_1)
# print(the_1)

lss = dataloader.get_agent_heading_list_byId(76,0,90)
print(lss)

lsls = dataloader.get_ego_vehicle_heading_list(0,90)
print(lsls)
# back_heading_list = dataloader.get_agent_heading_list_byId(back_car_list[0],0,90)
# print(back_heading_list)
# print(agents_list[math_tool.Id_to_Index(agents_list,1027)].states_mask_list)
# q=0
# for i in range(0,90):
#     v1,v2 = math_tool.find_front_back_vehicle(scenario, i,1027)
#     if v2 == 1033:
#         q += 1
#     print((v1,v2))
# print(q)
# print(v1.states_list[scenario.current_time_index].agent_id)


# l = math_tool.on_which_lane(scenario,agent_x,agent_y,agent_length,agent_width)
# print(l)


# l1 = math_tool.on_which_lane(scenario,v1.states_list[scenario.current_time_index].center_x,v1.states_list[scenario.current_time_index].center_y,v1.states_list[scenario.current_time_index].length,v1.states_list[scenario.current_time_index].width)

# print(l1)
# print(v1.states_array[scenario.current_time_index][3])
# print(v1.states_array[scenario.current_time_index][4])
# print(v1.states_array[scenario.current_time_index][5])
# print((v1.states_array[scenario.current_time_index][9],cur_state_sdc.heading ))

# l2 = math_tool.on_which_lane(scenario,v2.states_list[scenario.current_time_index].center_x,v2.states_list[scenario.current_time_index].center_y,v2.states_list[scenario.current_time_index].length,v2.states_list[scenario.current_time_index].width)

# print(l2)

# print((v2.states_array[scenario.current_time_index][9],cur_state_sdc.heading ))

# for time_index in time_index_list:
#     v1,v2 = math_tool.find_preceeding_following_in_lane(scenario, time_index, agent_id)
#     if v1 is not None:
#         print(time_index)
#         print("V1:")
#         print(v1.agent_id)
#         print((v1.states_list[time_index].center_x,v1.states_list[time_index].center_y))
#     if v2 is not None:
#         print("V2:")
#         print(v2.agent_id)  